/**
 * @file SolverNetAmb.cpp
 * Class to fix the undifferenced ambiguities and solve
 * the biass for satellite and stations simutaneously.
 *
 * 2019/09/01
 * major revision, add ambPrefixed to keep the consistency of the ambiguity datums
 *
 * 2019/10/06
 * major revision
 * introduce a sation-satellite-iteration method for phase-bias computation
 * shjzhang 2019/07/01
 *
 * 2019/10/06
 * remove Class "Ambiguity", replace it with "Variable"
 * shjzhang 2019/07/01
 *
 * 2020/01/26
 * modify all codes to support C11.
 *
 * 2020/01/27
 * polish the code and algorithm
 * 1. choose independent datum
 * 2. solve phase bias (upds) for sats and sources
 * 3. according the to the previous determined sat/source biases(upds)
 *    fix all the potential ambiguities
 * 4. update sat/source biases.
 * 5. store biases for next epoch computing.
 *
 */

#include "Edge.hpp"
#include "SolverNetAmb.hpp"
#include "MatrixFunctors.hpp"
#include "Matrix.hpp"
#include "Kruskal.hpp"
#include "Counter.hpp"

#define debug1 0
#define debug2 0

namespace gpstk
{

      // Returns a string identifying this object.
    std::string SolverNetAmb::getClassName() const
    { return "SolverNetAmb"; }


    void SolverNetAmb::init()
    {};

    // VariableEstDataMap interface
    int SolverNetAmb::Solve(const CommonTime& epoch, 
                            VariableEstDataMap& varEstData)
    {

        // clear satBiasData and sourceBiasData firstly
        satBiasData.clear();
        sourceBiasData.clear();

        // solver initial bias values at the first time
        if(firstTime)
        {
            // create edge
            // edge = variable(source, sat, arc) + variance
            EdgeSet currentEdgeSet;
            for(auto ved: varEstData)
            {
                Edge edge( ved.first, ved.second.variance);
                currentEdgeSet.insert(edge);
            }

            // independent edges
            Kruskal kruskal;
            kruskal.setEdges(currentEdgeSet);

            // create the mst
            EdgeSet indepEdgeSet;
            indepEdgeSet = kruskal.createMST();

            // Get independent ambiguity bias observations
            VariableEstDataMap indepAmbBiasData;
            for(auto edge: indepEdgeSet)
            {
                Variable ambVar( edge );

                double ambFloat, ambFixed, variance;

                ambFloat = varEstData[ambVar].value; 
                ambFixed = std::floor( ambFloat + 0.5);
                variance = varEstData[ambVar].variance;

                indepAmbBiasData[ambVar].value = ambFloat - ambFixed;
                indepAmbBiasData[ambVar].variance = variance;
            }

            // solve biass
            solveBiasEquations(indepAmbBiasData);

/*
            // store current satBiasData/sourceBiasData 
            oldSatBiasData = satBiasData;
            oldSourceBiasData = sourceBiasData;
*/
            for(auto sbd: satBiasData)
            {
                oldSatBiasData[sbd.first] = sbd.second;
            }

            for(auto sbd: sourceBiasData )
            {
                oldSourceBiasData[sbd.first] = sbd.second;
            }

            // no longer firstTime
            firstTime = false;
        }


        // fix all the other ambiguities
        ambEstDataFixed.clear();
        VariableEstDataMap ambBiasData;
        SatIDSet newSatSet;
        SourceIDSet newSourceSet;
        for(auto ved: varEstData)
        {
            SatID sat = ved.first.getSatellite();
            SourceID source = ved.first.getSource();

            bool findSat(true); 
            bool findSource(true);
            if( oldSatBiasData.find(sat) == oldSatBiasData.end())
            {
                findSat = false;
                newSatSet.insert(sat);
            }
            if( oldSourceBiasData.find(source) == oldSourceBiasData.end() )
            {
                findSource = false;
                newSourceSet.insert(source);
            }

            // only we have biases for both current sat and source 
            if( findSat && findSource )
            {
                double satBias = oldSatBiasData[sat];
                double sourceBias = oldSourceBiasData[source];
                double ambFloat = ved.second.value;
                double ambInt = ambFloat - sourceBias - satBias ;
                double ambVariance = ved.second.variance;

                ARData arData = ambRound.resolve(ambInt, ambVariance );
                if(arData.fixed)
                {
                    // bias value
                    double ambBias = ambFloat - arData.ambFixed;

                    // biasData 
                    EstData biasData;
                    biasData.value = ambBias;
                    biasData.variance = ambVariance;

                    // insert data into ambBiasData for bias estimation
                    ambBiasData[ved.first] = biasData;

                    // ambData 
                    EstData ambData;
                    ambData.value = arData.ambFixed;
                    ambData.variance = ambVariance;

                    // ambEstDataFixed
                    ambEstDataFixed[ved.first] = ambData;
                }
            }
        }

        // create edge
        // edge = variable(source, sat, arc) + variance
        EdgeSet currentEdgeSet;
        for(auto ved: ambBiasData)
        {
            Edge edge( ved.first, ved.second.variance);
            currentEdgeSet.insert(edge);
        }

        // independent edges
        Kruskal kruskal;
        kruskal.setEdges(currentEdgeSet);

        // create the mst
        EdgeSet indepEdgeSet;
        indepEdgeSet = kruskal.createMST();

        SourceIDSet validSourceSet = kruskal.getConnectedSourceSat();
        SatIDSet validSatSet = kruskal.getConnectedSatSat();

        VariableEstDataMap tempAmbBiasData;

        for(auto abd: ambBiasData)
        {
            SourceID source = abd.first.getSource();
            SatID sat = abd.first.getSatellite();

            bool isSourceValid(false);
            bool isSatValid(false);
            if(validSourceSet.find(source) != validSourceSet.end())
            {
                isSourceValid = true;
            }

            if(validSatSet.find(sat) != validSatSet.end())
            {
                isSatValid = true;
            }

            if(isSourceValid && isSatValid)
            {
                tempAmbBiasData[abd.first] = abd.second;
            }
        }

        // update bias observables
        ambBiasData = tempAmbBiasData;

        // solve biass again using all fixed ambiguities
        solveBiasEquations(ambBiasData);

        // compute biases for new satellites
        if(newSatSet.size() > 0)
        {
            for(auto newSat: newSatSet )
            {
                YDSTime yds(epoch);

                std::vector<double> biasVec;
                VariableEstDataMap ambEstDataNewSat;
                for(auto ved: varEstData)
                {
                    SatID varSat = ved.first.getSatellite();
                    if( varSat == newSat)
                    {

                        SourceID varSource = ved.first.getSource();
                        if( sourceBiasData.find(varSource) != sourceBiasData.end())
                        {
                            double sourceBias = sourceBiasData[varSource];
                            double ambFloat = ved.second.value - sourceBias;
                            double ambInt = std::floor(ambFloat + 0.5);
                            double satBias = ambFloat - ambInt;

                            biasVec.push_back(satBias);

                            // store the float amb value/variance for current sat
                            ambEstDataNewSat[ved.first] = ved.second; 
                        }
                    }
                }

                // compute the robust bias for current sat.
                double satBias = median<double>(biasVec);


                VariableEstDataMap ambEstDataNewSatFixed;

                // update satBias
                double bias(0.0);
                for(auto aed: ambEstDataNewSat)
                {
                    SourceID ambSource = aed.first.getSource();
                    double sourceBias = sourceBiasData[ambSource];
                    double ambFloat = aed.second.value;


                    double ambInt = ambFloat - sourceBias - satBias; 
                    double ambVar   = aed.second.variance;

                    ARData arData = ambRound.resolve(ambInt, ambVar);;

                    if(arData.fixed)
                    {
                        // ambData 
                        EstData ambData;
                        ambData.value = arData.ambFixed;
                        ambData.variance = ambVar;

                        // store  
                        ambEstDataNewSatFixed[aed.first] = ambData;
                    }
                }

                std::vector<double> biasVec2;
                for(auto aed: ambEstDataNewSatFixed)
                {
                    SourceID ambSource = aed.first.getSource();
                    double sourceBias = sourceBiasData[ambSource];
                    double ambFloat = ambEstDataNewSat[aed.first].value;
                    double ambFixed = ambEstDataNewSatFixed[aed.first].value;


                    double satBias = ambFloat - ambFixed - sourceBias;

                    biasVec2.push_back(satBias);
                }

                Stats<double> stats;
                stats.Add(biasVec2);

                satBiasData[newSat] = stats.Average();
            }
        }

        // bias connect between epochs
        biasConnect();       
        biasShift();

/*
        // store new sat/source biases
        oldSatBiasData = satBiasData;
        oldSourceBiasData = sourceBiasData;
*/

        for(auto sbd: satBiasData)
        {
            oldSatBiasData[sbd.first] = sbd.second;
        }

        for(auto sbd: sourceBiasData )
        {
            oldSourceBiasData[sbd.first] = sbd.second;
        }

    }


    // solve the bias equations as whole
    int SolverNetAmb::solveBiasEquations( VariableEstDataMap& ambBiasData )
    {
        // first check the sourceSet and satSet between input the the observations

        SourceIDSet sourceSet;
        SatIDSet satSet;
        for(auto var: ambBiasData)
        {
            satSet.insert(var.first.getSatellite());
            sourceSet.insert(var.first.getSource());
        }

//      for(auto var: ambBiasData)
//      {
//          cout << "var:" << var.first << " value:" << var.second.value << endl;
//      }

        // now, first, let's define the variables.
        Variable biasSta;
        biasSta.setType(TypeID::biasWL12);
        biasSta.setSourceIndexed(true);

        Variable biasSat;
        biasSat.setType(TypeID::biasSatWL12);
        biasSat.setSourceIndexed(false); // by fault, is source-indexed
        biasSat.setSatIndexed(true);

        // Prepare unknowns for current bias equations
        VariableSet unkSet;
        VariableSet sourceUnkSet, satUnkSet;
        for(auto source: sourceSet)
        {
            Variable var(biasSta);
            var.setSource( source );

            unkSet.insert( var );
            sourceUnkSet.insert(var);
        }

        for(auto sat: satSet)
        {
            Variable var(biasSat);
            var.setSatellite( sat  );
            unkSet.insert(var);
            satUnkSet.insert(var);
        }

           // Total number of the current equations
           // numEqu = numObs + 1; // 1 constraint equation
        int numEqu( ambBiasData.size() + 1 );
        int numVar( unkSet.size() );

        if(numVar==0)
        {
           GPSTK_THROW(InvalidSolver("bias var number is empty, you must set it first!"));
           exit(-1);
        }
        
        Matrix<double> hMatrix, rMatrix;
           // Resize hMatrix and rMatrix
        hMatrix.resize( numEqu, numVar, 0.0);
        rMatrix.resize( numEqu, numEqu, 0.0);

           // Declare temporal storage for values
        Vector<double> prefit(numEqu,0.0);

        int row(0);
        for(auto abd: ambBiasData)
        {
              // ambiguity
            Variable ambVar( abd.first );

            double tempBiasObs( abd.second.value );
            double tempBiasObsVar( abd.second.variance);

              // prefit
            prefit(row) = tempBiasObs;

              // rMatrix
              // convert variance to weight
            if(tempBiasObsVar<0.0001)
            {
                rMatrix(row, row) = 1.0/0.0001;
            }
            else
            {
                rMatrix(row, row) = 1.0/tempBiasObsVar;
            }

              // now, hMatrix

            Variable tempStaVar(biasSta);
            tempStaVar.setSource( ambVar.getSource() );

               // Now, Let's get the position of then tempStaVar in 'unkSet'
            int colSta(0);
            auto itVar=unkSet.begin();
            while( (*itVar) != tempStaVar )
            {
               colSta++;
               itVar++;
            }

               // Set the geometry matrix
            hMatrix(row, colSta) = 1.0;

            Variable tempSatVar(biasSat);
            tempSatVar.setSatellite( ambVar.getSatellite() );

               // Now, Let's get the position of then tempStaVar in 'unkSet'
            int colSat(0);
            auto itVar2=unkSet.begin();
            while( (*itVar2) != tempSatVar )
            {
               colSat++;
               itVar2++;
            }

               // Set the geometry matrix
            hMatrix(row, colSat) = 1.0;

            row++;

        }

        //=======================================================
        // constraint
        //=======================================================

          // very big weight
        rMatrix(row, row) = 1.0E+8;

          // now, hMatrix
        Variable tempSatVar(biasSat);
        datumSat = (*satSet.begin());
        tempSatVar.setSatellite( datumSat);
        prefit(row) = 0.0;

           // Now, Let's get the position of the datum sat
        int colDatum(0);
        VariableSet::const_iterator itVar=unkSet.begin();
        while( (*itVar) != tempSatVar )
        {
           colDatum++;
           itVar++;
        }
        hMatrix(row, colDatum) = 1.0;



        //=======================================================
        // end of constraint
        //=======================================================

        // now, compute nMatrix; wVector
        Matrix<double> nMatrix;
        Vector<double> wVector;

        Matrix<double> hT = transpose(hMatrix);

        nMatrix = hT * rMatrix * hMatrix;
        wVector = hT * rMatrix * prefit;

        int numSrc(sourceUnkSet.size());
        int numSat(satUnkSet.size());

        // now, extract source-related Normal equation from nMatrix
        Matrix<double> Nrr(numSrc, numSrc, 0.0);
        Matrix<double> Nrs(numSrc, numSat, 0.0);
        Matrix<double> Nss(numSat, numSat, 0.0);
        Matrix<double> Nsr(numSat, numSrc, 0.0);
        Vector<double> Wr(numSrc,0.0), Ws(numSat,0.0); 

        // Nrr
        for(int i=0;i<numSrc;i++)
        {
            for(int j=0; j<numSrc;j++)
            {
                Nrr(i,j) = nMatrix(i,j);
            }
        }

        // Nrs
        for(int i=0;i<numSrc;i++)
        {
            for(int j=numSrc; j<numVar; j++)
            {
                Nrs(i,j-numSrc) = nMatrix(i,j);
            }
        }

        // Nsr
        Nsr = transpose(Nrs);

        // Nss
        for(int i=numSrc;i<numVar;i++)
        {
            for(int j=numSrc; j<numVar; j++)
            {
                Nss(i-numSrc,j-numSrc) = nMatrix(i,j);
            }
        }

        for(int i=0;i<numSrc;i++)
        {
            Wr(i) = wVector(i);
        }

        for(int i=numSrc;i<numVar;i++)
        {
            Ws(i-numSrc) = wVector(i);
        }


        if(!Nrr.isDiagonal())
        {
            MatrixException e("Nrr is not diagnal!");
            cerr << e << endl; 
            GPSTK_THROW(e);
        }
        if(!Nss.isDiagonal())
        {
            MatrixException e("Nss is not diagnal!");
            cerr << e << endl; 
            GPSTK_THROW(e);
        }

        // according to the block-wise least square method
        // if the parameters are divided into Xr and Xs, 
        // then,
        // Xs = inv[ Nss - Nsr*inv(Nrr)*Nrs ] * [ Ws - Nsr*inv(Nrr)*Wr ]
        // afterwards,
        // Xr = inv(Nrr) * (Wr - Nrs*Xs)

        Vector<double> Xx(numVar, 0.0);
        Vector<double> Xr(numSrc, 0.0);
        Vector<double> Xs(numSat, 0.0);

        if(usingBlockWiseLSQ)
        {
            Matrix<double> Nss2;
            Vector<double> Ws2;

            Matrix<double> invNrr(numSrc, numSrc, 0.0);

            double clock1 ( Counter::now() );

              // inverse of the Nrr
              // because the Nrr is diagnal, the inverse of Nrr
              // is equal with the inverse of the diagnal element of Nrr
            for(int i=0;i<numSrc;i++)
            {
                invNrr(i,i) = 1.0/Nrr(i,i);
            }

            Nss2 = Nss - Nsr * invNrr * Nrs;
            Ws2 = Ws - Nsr * invNrr * Wr;

            Matrix<double> invNss2;

            try
            {
                invNss2 = inverseChol(Nss2);
            }
            catch(Exception& e)
            {
                cout << "Nss:" << endl;
                cout << Nss << endl;
                cout << "Nrr" << endl;
                cout << Nrr << endl;
                cout << "Nss2" << endl;
                cout << Nss2 << endl;

                InvalidSolver is("solveBiasEquations(): Unable to compute Nss2.");
                cerr << is << endl;
                exit(-1);
            }

            Xs = invNss2 * Ws2;
            Xr = invNrr * ( Wr - Nrs*Xs);

            double clock2 ( Counter::now() );

        }
        else
        {
            double clock1 ( Counter::now() );
            // inverse directly
            Xx = inverseChol( nMatrix ) * wVector;
            // Xr
            for(int i=0;i<numSrc;i++)
            {
                Xr(i) = Xx(i);
            }
            // Xs
            for(int i=numSrc;i<numVar;i++)
            {
                Xs(i-numSrc) = Xx(i);
            }

            double clock2 ( Counter::now() );
        }

        sourceBiasData.clear();
        row = 0;
        for(SourceIDSet::const_iterator srcIt=sourceSet.begin();
            srcIt!=sourceSet.end();
            ++srcIt)
        {
            sourceBiasData[(*srcIt)] = Xr(row);
            row++;

        }

        satBiasData.clear();
        row = 0;
        for(SatIDSet::const_iterator satIt=satSet.begin();
            satIt!=satSet.end();
            ++satIt)
        {
            satBiasData[(*satIt)] = Xs(row);
            row++;

        }
    }

    void SolverNetAmb::biasConnect()
    {
        for(auto sbd: satBiasData)
        {
            if(oldSatBiasData.find(sbd.first)!= oldSatBiasData.end())
            {
                double biasCurrent = sbd.second;
                double biasOld = oldSatBiasData[ sbd.first ];

                double diff = biasCurrent - biasOld;
                biasCurrent = biasCurrent - std::floor( diff + 0.5);
                satBiasData[sbd.first] = biasCurrent;
            }
        }

        for(auto sbd: sourceBiasData)
        {
            if(oldSourceBiasData.find(sbd.first)!=oldSourceBiasData.end())
            {
                double biasCurrent = sbd.second;
                double biasOld = oldSourceBiasData[ sbd.first ];

                double diff = biasCurrent - biasOld;
                biasCurrent = biasCurrent - std::floor( diff + 0.5);
                sourceBiasData[sbd.first] = biasCurrent;
            }
        }
    };

    void SolverNetAmb::biasShift()
    {
        // compute bias
        double meanOld(0.0), meanCurrent(0.0);
        int numSat(0);
        for(auto sbd: satBiasData)
        {
            if(oldSatBiasData.find(sbd.first) != oldSatBiasData.end())
            {
                meanCurrent += sbd.second;
                meanOld += oldSatBiasData[sbd.first];
                numSat++;
            }
        }

        double bias;
        bias = ( meanCurrent - meanOld )/numSat; 
        
        // shift the systematic bias for all satellite biases
        for(auto sbd: satBiasData)
        {
            satBiasData[sbd.first] -= bias;
        }

        // compute source bias
        meanOld = 0.0;
        meanCurrent = 0.0;
        int numSrc(0);
        for(auto sbd: sourceBiasData)
        {
            if(oldSourceBiasData.find(sbd.first) != oldSourceBiasData.end())
            {
                meanCurrent += sbd.second;
                meanOld += oldSourceBiasData[sbd.first];
                numSrc++;
            }
        }

        bias = ( meanCurrent - meanOld )/numSrc; 
        
        // shift the systematic bias for all satellite biases
        for(auto sbd: sourceBiasData)
        {
            sourceBiasData[sbd.first] -= bias;
        }
    };


    satValueMap SolverNetAmb::getSatBiasData() const
    {
        return satBiasData;
    }

    sourceValueMap SolverNetAmb::getSourceBiasData() const
    {
        return sourceBiasData;
    }

    VariableEstDataMap SolverNetAmb::getFixedAmbEstData() const
    {
        return ambEstDataFixed;
    }

}  // End of namespace gpstk
